#!/usr/bin/python2.4
#
# Copyright 2008 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Contains the methods to import mail via Google Apps Email Migration API.

  MigrationService: Provides methods to import mail.
"""

__author__ = ('google-apps-apis@googlegroups.com',
              'pti@google.com (Prashant Tiwari)')


import base64
import threading
import time
from atom.service import deprecation
from gdata.apps import migration
from gdata.apps.migration import MailEntryProperties
import gdata.apps.service
import gdata.service


API_VER = '2.0'


class MigrationService(gdata.apps.service.AppsService):
    """Client for the EMAPI migration service.  Use either ImportMail to import
    one message at a time, or AddMailEntry and ImportMultipleMails to import a
    bunch of messages at a time.
    """

    def __init__(self, email=None, password=None, domain=None, source=None,
                 server='apps-apis.google.com', additional_headers=None):
        gdata.apps.service.AppsService.__init__(
            self, email=email, password=password, domain=domain, source=source,
            server=server, additional_headers=additional_headers)
        self.mail_batch = migration.BatchMailEventFeed()
        self.mail_entries = []
        self.exceptions = 0

    def _BaseURL(self):
        return '/a/feeds/migration/%s/%s' % (API_VER, self.domain)

    def ImportMail(self, user_name, mail_message, mail_item_properties,
                   mail_labels):
        """Imports a single mail message.

        Args:
          user_name: The username to import messages to.
          mail_message: An RFC822 format email message.
          mail_item_properties: A list of Gmail properties to apply to the message.
          mail_labels: A list of labels to apply to the message.

        Returns:
          A MailEntry representing the successfully imported message.

        Raises:
          AppsForYourDomainException: An error occurred importing the message.
        """
        uri = '%s/%s/mail' % (self._BaseURL(), user_name)

        mail_entry = migration.MailEntry()
        mail_entry.rfc822_msg = migration.Rfc822Msg(text=(base64.b64encode(
            mail_message)))
        mail_entry.rfc822_msg.encoding = 'base64'
        mail_entry.mail_item_property = map(
            lambda x: migration.MailItemProperty(value=x), mail_item_properties)
        mail_entry.label = map(lambda x: migration.Label(label_name=x),
                               mail_labels)

        try:
            return migration.MailEntryFromString(str(self.Post(mail_entry, uri)))
        except gdata.service.RequestError, e:
            # Store the number of failed imports when importing several at a time
            self.exceptions += 1
            raise gdata.apps.service.AppsForYourDomainException(e.args[0])

    def AddBatchEntry(self, mail_message, mail_item_properties,
                      mail_labels):
        """Adds a message to the current batch that you later will submit.

        Deprecated, use AddMailEntry instead

        Args:
          mail_message: An RFC822 format email message.
          mail_item_properties: A list of Gmail properties to apply to the message.
          mail_labels: A list of labels to apply to the message.

        Returns:
          The length of the MailEntry representing the message.
        """
        deprecation("calling deprecated method AddBatchEntry")
        mail_entry = migration.BatchMailEntry()
        mail_entry.rfc822_msg = migration.Rfc822Msg(text=(base64.b64encode(
            mail_message)))
        mail_entry.rfc822_msg.encoding = 'base64'
        mail_entry.mail_item_property = map(
            lambda x: migration.MailItemProperty(value=x), mail_item_properties)
        mail_entry.label = map(lambda x: migration.Label(label_name=x),
                               mail_labels)

        self.mail_batch.AddBatchEntry(mail_entry)

        return len(str(mail_entry))

    def SubmitBatch(self, user_name):
        """Sends all the mail items you have added to the batch to the server.

        Deprecated, use ImportMultipleMails instead

        Args:
          user_name: The username to import messages to.

        Returns:
          An HTTPResponse from the web service call.

        Raises:
          AppsForYourDomainException: An error occurred importing the batch.
        """
        deprecation("calling deprecated method SubmitBatch")
        uri = '%s/%s/mail/batch' % (self._BaseURL(), user_name)

        try:
            self.result = self.Post(self.mail_batch, uri,
                                    converter=migration.BatchMailEventFeedFromString)
        except gdata.service.RequestError, e:
            raise gdata.apps.service.AppsForYourDomainException(e.args[0])

        self.mail_batch = migration.BatchMailEventFeed()

        return self.result

    def AddMailEntry(self, mail_message, mail_item_properties=None,
                     mail_labels=None, identifier=None):
        """Prepares a list of mail messages to import using ImportMultipleMails.

        Args:
          mail_message: An RFC822 format email message as a string.
          mail_item_properties: List of Gmail properties to apply to the
              message.
          mail_labels: List of Gmail labels to apply to the message.
          identifier: The optional file identifier string

        Returns:
          The number of email messages to be imported.
        """
        mail_entry_properties = MailEntryProperties(
            mail_message=mail_message,
            mail_item_properties=mail_item_properties,
            mail_labels=mail_labels,
            identifier=identifier)

        self.mail_entries.append(mail_entry_properties)
        return len(self.mail_entries)

    def ImportMultipleMails(self, user_name, threads_per_batch=20):
        """Launches separate threads to import every message added by AddMailEntry.

        Args:
          user_name: The user account name to import messages to.
          threads_per_batch: Number of messages to import at a time.

        Returns:
          The number of email messages that were successfully migrated.

        Raises:
          Exception: An error occurred while importing mails.
        """
        num_entries = len(self.mail_entries)

        if not num_entries:
            return 0

        threads = []
        for mail_entry_properties in self.mail_entries:
            t = threading.Thread(name=mail_entry_properties.identifier,
                                 target=self.ImportMail,
                                 args=(user_name, mail_entry_properties.mail_message,
                                       mail_entry_properties.mail_item_properties,
                                       mail_entry_properties.mail_labels))
            threads.append(t)
        try:
            # Determine the number of batches needed with threads_per_batch in each
            batches = num_entries / threads_per_batch + (
                0 if num_entries % threads_per_batch == 0 else 1)
            batch_min = 0
            # Start the threads, one batch at a time
            for batch in range(batches):
                batch_max = ((batch + 1) * threads_per_batch
                             if (batch + 1) * threads_per_batch < num_entries
                             else num_entries)
                for i in range(batch_min, batch_max):
                    threads[i].start()
                    time.sleep(1)

                for i in range(batch_min, batch_max):
                    threads[i].join()

                batch_min = batch_max

            self.mail_entries = []
        except Exception, e:
            raise Exception(e.args[0])
        else:
            return num_entries - self.exceptions
